"""
设备配置模块 - 自动检测并配置 CUDA/CPU
"""
import torch

# 自动检测 CUDA 是否可用
CUDA_AVAILABLE = torch.cuda.is_available()
DEVICE = torch.device('cuda' if CUDA_AVAILABLE else 'cpu')

def get_device():
    """返回当前设备 (cuda 或 cpu)"""
    return DEVICE

def to_device(tensor_or_module):
    """将 tensor 或 module 移动到可用设备"""
    return tensor_or_module.to(DEVICE)

def print_device_info():
    """打印设备信息"""
    if CUDA_AVAILABLE:
        print(f"✓ 使用 CUDA (GPU): {torch.cuda.get_device_name(0)}")
    else:
        print("⚠ CUDA 不可用，使用 CPU")
        print("提示：如需使用 GPU，请安装 CUDA 版本的 PyTorch")

